import numpy as np
from tensorflow.keras.models import load_model
import cv2

# 种类字典
class_dict = {0: '猫', 1: '狗'}


def predict(img_path):
    # 载入模型
    model = load_model('./model4.h5')
    # 载入图片，并处理
    img = cv2.imread(img_path)
    img = cv2.resize(img, (64, 64))
    img_RGB = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img_nor = img_RGB / 255
    img_nor = np.expand_dims(img_nor, axis=0)

    # 预测
    # print((np.argmax(model.predict(img_nor))))
    # print(model.predict(img_nor))
    y = model.predict_classes(img_nor)
    print(class_dict.get(y[0][0]))  # 直接输出种类 0是猫 1是狗


if __name__ == "__main__":
    predict('./datasets/cats_and_dogs/test/cat.1500.jpg')  # 0
    predict('./datasets/cats_and_dogs/test/dog.1500.jpg')  # 1
    predict('./datasets/cats_and_dogs/test/dog.1504.jpg')  # 1
